/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * License); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * AS IS BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

/*
 * Copyright (c) 2019, Open AI Lab
 * Author: jingyou@openailab.com
 */

#ifndef __REF_REORG_KERNEL_H__
#define __REF_REORG_KERNEL_H__

#include <stdint.h>
#include <math.h>

#include "compiler_fp16.h"

#ifdef __cplusplus
extern "C" {
#endif

struct ref_reorg_param
{
    int w;
    int h;
    int c;
    int batch;
    int stride;
};

typedef int (*ref_reorg_kernel_t)(const void* in_data, void* out_data, ref_reorg_param* param);

static int ref_reorg_common(const char* in_data, char* out_data, ref_reorg_param* param, int element_size)
{
    int w = param->w;
    int h = param->h;
    int c = param->c;
    int batch = param->batch;
    int stride = param->stride;

    int out_c = c / (stride * stride);

    for(int b = 0; b < batch; ++b)
    {
        for(int k = 0; k < c; ++k)
        {
            for(int j = 0; j < h; ++j)
            {
                for(int i = 0; i < w; ++i)
                {
                    int in_index = i + w * (j + h * (k + c * b));
                    int c2 = k % out_c;
                    int offset = k / out_c;
                    int w2 = i * stride + offset % stride;
                    int h2 = j * stride + offset / stride;
                    int out_index = w2 + w * stride * (h2 + h * stride * (c2 + out_c * b));
                    memcpy(out_data + in_index * element_size, in_data + out_index * element_size, element_size);
                }
            }
        }
    }
    return 0;
}

#ifdef CONFIG_KERNEL_FP32
#include "ref_reorg_fp32.c"
#endif

#ifdef CONFIG_KERNEL_FP16
#include "ref_reorg_fp16.c"
#endif

#ifdef CONFIG_KERNEL_INT8
#include "ref_reorg_int8.c"
#endif

#ifdef CONFIG_KERNEL_UINT8
#include "ref_reorg_uint8.c"
#endif

#ifdef __cplusplus
}
#endif

#endif
